
% Strain from Niagara MITgcm runs. The techique is based on Kunze et al
% 2006 which is strain = [N^2-\bar{N}^2/(\bar{N}^2)]
%
% This is to work with Niagara MITgcm runs in the selected depth range
%
% Ritabrata Thakur 05 Nov 2021. Completed 23 Nov 2021
%
% Last used: 31 Mar 2022

clear

%% This is an automatic routine that saves the spectra of strain for
% all the vertical level runs, for all the locations at MMP sites, and for
% both background on and off. 

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
run_no_array = {19,20,1,2,3,4};
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

for jj=1:length(run_no_array)

run_no = run_no_array{jj};

%check for "run_no" variable in the excel sheet
if (run_no == 19 || run_no == 20)
 vert_levs = 109;
end
if (run_no == 1 || run_no == 2)
    vert_levs = 153;
end
if (run_no == 3 || run_no == 4)
    vert_levs = 264;
end

%if (run_no == 34 || run_no == 35)
%    vert_levs = 153;
%end

%if (run_no == 25 || run_no == 26)
%    vert_levs = 264;
%end

fprintf('Running vertical level %d\n',vert_levs);
for MMP_number=[1 3 4]
fprintf('Running MMP_number %d\n',MMP_number);
clearvars -except jj run_no vert_levs MMP_number run_no_array 

%% Reading variables
if (run_no == 19 || run_no == 20)
 [data,lat,lon,ix,jx,ncpath,XC,YC,RC,time_start,time_end,time_array,z_ip] = read_ncfiles(vert_levs,MMP_number,run_no);
end

if (run_no == 1 || run_no == 2 || run_no == 3 || run_no == 4)
 [data,lat,lon,ix,jx,ncpath,XC,YC,RC,time_start,time_end,time_array,z_ip] = read_ncfiles_km_runs(vert_levs,MMP_number,run_no);
end

% the inertial periods would be required if I separate the strain into
% different components
if (MMP_number == 1)
 inertial_start = 25.05; % 90% of 27.84 hrs
 inertial_end   = 30.6;  % 110% of 27.84 hrs
end

if (MMP_number == 3)
 inertial_start = 22.34; % 90% of 24.83 hrs
 inertial_end   = 27.31; % 110% of 24.83 hrs
end

if (MMP_number == 4)
 inertial_start = 21.52; % 90% of 23.92 hrs
 inertial_end   = 26.31; % 110% of 23.92 hrs
end

%% N2 calculation for strain

%Z = depth_array;
p = gsw_p_from_z(squeeze(RC),lat);

% Taking salt and salinity from the "all_ON" run
[SA, ~] = gsw_SA_from_SP(squeeze(data.salt(1,1,:,:)),p,lon,lat);  %converts practical salinity
%to absolute salinity which is used by the N2 routine. The flag of "1" if
%the measurement is in ocean. If it is used in land, set it to 0. do not
%put anything. it is automatically set
CT = gsw_CT_from_t(SA,squeeze(data.temp(1,1,:,:)),p); %converts in-situ temperature to
%conservative temperature as is required in GSW toolbox

clear N2 NN
for i=1:length(CT(1,:))
 [N2,p_mid] = gsw_Nsquared(SA(:,i),CT(:,i),p,lat);
 NN(i,:)=N2;  %this is the total Brunt-vaisala frequency squared stored
 % at a single location. It is probably best to continue with one location
 % as averaging over different locations would mean we would have to
 % calculate the wkb factor for each depth differently as the
 % stratification would vary according to location
end

% removing the values where there is unstable stratification
% because when we calculate Brunt Vaisala frequency, it would otherwise
% become complex
for i=1:length(NN)
    for j=1:length(NN(1,:))
      if(NN(i,j) <= 0)
      NN(i,j) = 0;
      end
    end
end

% choosing intermediate depth range of MMP range. This can be used for both 109 level
% and 264 level. There might be slight differences but that should be okay
start_depth=find(abs(p_mid-80)<5,1); %taking approx 80m as the start depth to match MMP 
end_depth=find(abs(p_mid-1400)<50,1,'last');  %taking approx 1400m as the end depth to match MMP 
% Brunt Vaisala in the intermediate depth range to compare with MMP

%sqrt of the N^2
BV = sqrt(NN);

%% Strain calculation

clear BV_interp 
%interpolating to original depth levels. Remember that "p" and the
%depth_array are different 
for i=1:length(NN)
 BV_interp(i,:) = interp1(p_mid, BV(i,:), p, 'pchip');
end

% Brunt Vaisala in the intermediate depth range (approz 80--1400m)
depth_array_short = squeeze(RC(start_depth:end_depth));
BV_interp_short   = squeeze(BV_interp(time_start:time_end,start_depth:end_depth));
BV_interp_time_mean =  mean(BV_interp_short, 'omitnan');
BV_interp_time_mean = naninterp(BV_interp_time_mean);

NN_interp_short   = squeeze(NN(time_start:time_end,start_depth:end_depth));
NN_interp_time_mean =  mean(NN_interp_short, 'omitnan');
NN_interp_time_mean = naninterp(NN_interp_time_mean);

% figure
% plot(log10(abs(NN_interp_time_mean)))

clear strain
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Strain calculation
for i=1:length(depth_array_short)
  strain(:,i) = (NN_interp_short(:,i) - NN_interp_time_mean(i))./NN_interp_time_mean(i);
end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%

for i=1:length(strain(1,:))
  %strain_3_to_5hr(:,i)  = bandpass(strain(:,i),[1/5.20 1/2.80], 1); 
  %strain_6_to_8hr(:,i)  = bandpass(strain(:,i),[1/8.20 1/5.80], 1); 
  strain_highpass(:,i)  = highpass(strain(:,i),1/11.5, 1); 
  strain_lowpass(:,i)   = lowpass(strain(:,i),1/11.5, 1); 
  strain_sem_diur(:,i)  = bandpass(strain(:,i),[1/13.5 1/11.5], 1);   
  strain_inertial(:,i)  = bandpass(strain(:,i),[1/inertial_end 1/inertial_start], 1); 
end


%% Stretched depth calculation
% stretched depth from the chosen intermediate depth
clear z_stretched
N_0_mean = 0.0052; %used in literature = 3cph
z_stretched(1) = (-depth_array_short(1));
for ii=2:length(depth_array_short)
  z_stretched(ii) = z_stretched(1) + trapz(-depth_array_short(1:ii), squeeze(BV_interp_time_mean(1:ii)))./N_0_mean; 
end

z_stretched_equal_spaced = linspace(z_stretched(1), z_stretched(end), length(z_stretched(1:end)));
%delta_z = 2*pi/(abs(z_stretched_equal_spaced(2))-abs(z_stretched_equal_spaced(1))); 
%delta_z_unstretched = 2*pi/(abs(z_unstretched_equal(2)-z_unstretched_equal(1)));
% removed the 2pi in the wavenumber following early literature (Gregg et
% al 1991: "varieties of fully resolved spectra of vertical shear"): 4Nov21
delta_z = 1/(abs(z_stretched_equal_spaced(2))-abs(z_stretched_equal_spaced(1))); 
% to have it in cycles per km, multiply by thousand
NFFT_waveno_strain=length(squeeze(strain(1,:))); %taking FFT along the whole vertical grid
% These are the wavenumbers.  we can do this because the depth
% array is equally spaced albeit in stretched coordinates. This is
% different when we calculate both freq-waveno because there we create
% wavenos from -ve to +ve delta_z/2. 

waveno_strain = 0:delta_z/NFFT_waveno_strain:delta_z/2;

%% Strain PSD calculation
clear wave_psd_avg_strain wave_psd_avg_strain_3_to_5 wave_psd_avg_strain_6_to_8 wave_psd_avg_strain_inertial psdx_array_strain
%%%%%%%%%%%%%%%%
%[psdx_array_strain_3_to_5,wave_psd_avg_strain_3_to_5]     = waveno_only_spectra_strain_ybp(strain_3_to_5hr,z_stretched_equal_spaced,NFFT_waveno_strain);
%[psdx_array_strain_6_to_8,wave_psd_avg_strain_6_to_8]     = waveno_only_spectra_strain_ybp(strain_6_to_8hr,z_stretched_equal_spaced,NFFT_waveno_strain);
%[psdx_array_strain_inertial,wave_psd_avg_strain_inertial] = waveno_only_spectra_strain_ybp(strain_inertial,z_stretched_equal_spaced,NFFT_waveno_strain);
[psdx_array_strain,wave_psd_avg_strain]                   = waveno_only_spectra_strain_ybp(strain,z_stretched_equal_spaced,NFFT_waveno_strain);
[~, wave_psd_avg_strain_inertial]                         = waveno_only_spectra_strain_ybp(strain_inertial,z_stretched_equal_spaced,NFFT_waveno_strain);
[~, wave_psd_avg_strain_highpass]                         = waveno_only_spectra_strain_ybp(strain_highpass,z_stretched_equal_spaced,NFFT_waveno_strain);
[~, wave_psd_avg_strain_lowpass]                          = waveno_only_spectra_strain_ybp(strain_lowpass,z_stretched_equal_spaced,NFFT_waveno_strain);
[~, wave_psd_avg_strain_semdiur]                          = waveno_only_spectra_strain_ybp(strain_sem_diur,z_stretched_equal_spaced,NFFT_waveno_strain);
%%%%%%%%%%%%%%%%

%wave_psd_avg_strain_3_to_5   = wave_psd_avg_strain_3_to_5/(NFFT_waveno_strain*delta_z);
%wave_psd_avg_strain_6_to_8   = wave_psd_avg_strain_6_to_8/(NFFT_waveno_strain*delta_z);
%wave_psd_avg_strain_inertial = wave_psd_avg_strain_inertial/(NFFT_waveno_strain*delta_z);
wave_psd_avg_strain          = wave_psd_avg_strain./(NFFT_waveno_strain*delta_z);
wave_psd_avg_strain_inertial = wave_psd_avg_strain_inertial./(NFFT_waveno_strain*delta_z);
wave_psd_avg_strain_highpass = wave_psd_avg_strain_highpass./(NFFT_waveno_strain*delta_z);
wave_psd_avg_strain_lowpass  = wave_psd_avg_strain_lowpass./(NFFT_waveno_strain*delta_z);
wave_psd_avg_strain_semdiur  = wave_psd_avg_strain_semdiur./(NFFT_waveno_strain*delta_z);

psd_strain.waveno     = waveno_strain;
psd_strain.tot_strain = wave_psd_avg_strain;
psd_strain.inertial   = wave_psd_avg_strain_inertial;
psd_strain.highpass   = wave_psd_avg_strain_highpass;
psd_strain.lowpass    = wave_psd_avg_strain_lowpass;
psd_strain.semdiur    = wave_psd_avg_strain_semdiur;
psd_strain.comment = 'Tot_strain is the psd of strain. The depth range is near to MMP depth range of 80-1400m';

end
end
